package net.howson.phil.datahiding;

import java.awt.image.BufferedImage;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.Random;

import javax.imageio.ImageIO;
import javax.imageio.ImageReader;
import javax.imageio.stream.FileImageInputStream;

import net.howson.phil.datahiding.scheme.DefaultSingleValueEncodingScheme;
import net.howson.phil.datahiding.scheme.SingleValueEncodingScheme;
import net.howson.phil.datahiding.stream.ComponentOutputStream;
import net.howson.phil.datahiding.stream.ComponentRecoveryStream;
import net.howson.phil.datahiding.stream.HeadedInputStream;
import net.howson.phil.datahiding.stream.NibbleInputStream;
import net.howson.phil.datahiding.stream.NibbleOutputStream;

public class DataHide {

	private static final String EXTRACTOPTION = "-x";
	private static final String KEY = "-key";
	private static final String SEED = "-seed";
	private static final String OUTPUT = "-output";
	private static final String INPUT = "-input";
	private final File inputFile;
	private final File outputFile;
	private final File seedFile;

	public DataHide(String inputFile, String outputFile, String seedFile) {
		this.inputFile = new File(inputFile);
		this.outputFile = new File(outputFile);
		if (seedFile != null) {
			this.seedFile = new File(seedFile);
		} else {
			this.seedFile = null;
		}
	}

	public static void main(String[] argv) throws IOException {
		String inputFile = null;
		String outputFile = null;
		String seedFile = null;
		int randomSeed = 0;
		boolean extract = false;

		for (int i = 0; i < argv.length; ++i) {
			if (INPUT.equals(argv[i])) {
				inputFile = argv[++i];
			} else if (OUTPUT.equals(argv[i])) {
				outputFile = argv[++i];
			} else if (SEED.equals(argv[i])) {
				seedFile = argv[++i];
			} else if (KEY.equals(argv[i])) {
				randomSeed = Integer.parseInt(argv[++i]);
			} else if (EXTRACTOPTION.equals(argv[i])) {
				extract = true;
			} else {
				usage();
				throw new IllegalArgumentException("Unknown argument : "
						+ argv[i]);
			}
		}
		if (inputFile == null) {
			usage();
			throw new IllegalArgumentException("Require an input file : "
					+ inputFile);
		}
		if (outputFile == null) {
			usage();
			throw new IllegalArgumentException("Require an output file : "
					+ outputFile);
		}
		if (seedFile == null) {
			usage();
			throw new IllegalArgumentException("Require a seed file : "
					+ outputFile);
		}

		SingleValueEncodingScheme scheme = new DefaultSingleValueEncodingScheme(
				randomSeed);
		DataHide app = new DataHide(inputFile, outputFile, seedFile);
		if (extract) {
			app.extract(scheme);
		} else {
			app.encode(scheme);
		}

	}

	private void encode(SingleValueEncodingScheme scheme) throws IOException {

		BufferedImage seedImg = loadImage(seedFile);

		int width = seedImg.getWidth(null);
		int height = seedImg.getHeight(null);

		int encodableSize = scheme.getNumberOfBytesEncodable(width, height);
		long inputSize = inputFile.length();
		int repititions = (int) Math.ceil(((double) inputSize)
				/ (double) encodableSize);

		BufferedImage outImg = new BufferedImage(width * repititions, height,
				BufferedImage.TYPE_3BYTE_BGR);

		InputStream is = new HeadedInputStream(inputSize, repititions,
				new BufferedInputStream(new FileInputStream(inputFile)));
		NibbleInputStream nis = new NibbleInputStream(is);
		InputStream seedInputStream = new ComponentRecoveryStream(seedImg, repititions);
		ComponentOutputStream out = new ComponentOutputStream(outImg,
				repititions);
		RandomStream random = new RandomStream(new Random());
		int i;
		int j;
		while ((i = nis.read()) != -1 && (j = seedInputStream.read()) != -1) {
			out.write(scheme.encode(i, j));
		}
		while ((i = random.read()) != -1 && (j = seedInputStream.read()) != -1) {
			out.write(scheme.encode(i, j));
		}

		nis.close();
		seedInputStream.close();

		String[] tmp = outputFile.getName().split("\\.");
		String outputExtension = tmp[tmp.length - 1];

		ImageIO.write(outImg, outputExtension, outputFile);

	}

	private void extract(SingleValueEncodingScheme scheme) throws IOException {
		BufferedImage inputImage = loadImage(inputFile);
		BufferedImage seedImage;
		if (seedFile != null) {
			seedImage = loadImage(seedFile);
		} else {
			seedImage = null;
		}
		ComponentRecoveryStream inImageStream = new ComponentRecoveryStream(
				inputImage, 1);
		ComponentRecoveryStream inSeedStream = new ComponentRecoveryStream(
				seedImage, Integer.MAX_VALUE);
		int i;
		int j;

		ByteArrayOutputStream bos = new ByteArrayOutputStream();
		NibbleOutputStream headerOut = new NibbleOutputStream(bos);

		int read = 0;
		while (read < 16 && (i = inImageStream.read()) != -1
				&& (j = inSeedStream.read()) != -1 ) {
			headerOut.write(scheme.decode(i, j));
			++read;
		}
		headerOut.close();

		ByteBuffer header = ByteBuffer.wrap(bos.toByteArray());
		header.position(0);
		int size = header.getInt();
		int reps = header.getInt();

		inImageStream.setRepititions(reps);

		read = 0;

		OutputStream out = new NibbleOutputStream(new BufferedOutputStream(
				new FileOutputStream(outputFile)));

		int componentSize = size*2;
		while ((i = inImageStream.read()) != -1
				&& (j = inSeedStream.read()) != -1 && read < componentSize) {
			out.write(scheme.decode(i, j));
			++read;
		}
		out.close();

	}

	private BufferedImage loadImage(File inputFile)
			throws FileNotFoundException, IOException {
		String[] tmp = inputFile.getName().split("\\.");
		String inputSuffix = tmp[tmp.length - 1];

		for (Iterator<ImageReader> it = ImageIO
				.getImageReadersBySuffix(inputSuffix); it.hasNext();) {
			ImageReader reader = it.next();
			FileImageInputStream fileImageInputStream = new FileImageInputStream(
					inputFile);
			reader.setInput(fileImageInputStream);

			try {
				BufferedImage image = reader.read(0);
				if (image != null) {
					return image;
				}
			} catch (Exception ignored) {

			}
		}
		throw new IllegalArgumentException("Could not load image: " + inputFile);
	}

	private static void usage() {
		System.out.println("usage:");
		System.out.println("\t" + EXTRACTOPTION);
		System.out.println("\t\tExtract from image");
		System.out.println("\t" + KEY);
		System.out.println("\t\tInteger key");
		System.out.println("\t" + SEED);
		System.out.println("\t\tSeed image");
		System.out.println("\t" + OUTPUT);
		System.out.println("\t\tOutput file");
		System.out.println("\t" + INPUT);
		System.out.println("\t\tInput file");

	}

	public static class RandomStream extends InputStream {

		private final Random random;

		public RandomStream(Random random) {
			this.random = random;
		}

		@Override
		public int read() throws IOException {
			return random.nextInt(16);
		}

	}

}